use crate::auth::Authenticator;
use crate::config::ServerConfig;
use crate::traffic::TrafficStats;
use hysteria_core::protocol::{AuthRequest, AuthResponse, TcpRequest, TcpResponse, URL_HOST, URL_PATH, STATUS_AUTH_OK};
use hysteria_core::crypto::SalamanderObfuscator;
use std::sync::Arc;
use std::net::SocketAddr;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::{error, info, warn, debug};
use quinn::{Endpoint, ServerConfig as QuinnServerConfig};
use rustls::ServerConfig as RustlsServerConfig;

pub struct HysteriaServer {
    config: Arc<ServerConfig>,
    traffic_stats: Arc<TrafficStats>,
    authenticator: Arc<dyn Authenticator>,
    obfuscator: Option<Arc<SalamanderObfuscator>>,
}

impl HysteriaServer {
    pub fn new(
        config: ServerConfig,
        authenticator: Arc<dyn Authenticator>,
    ) -> anyhow::Result<Self> {
        let traffic_stats = Arc::new(TrafficStats::new());
        
        // Initialize obfuscator if configured
        let obfuscator = if let Some(salamander_config) = &config.obfs.salamander {
            Some(Arc::new(SalamanderObfuscator::new(&salamander_config.password)))
        } else {
            None
        };

        Ok(Self {
            config: Arc::new(config),
            traffic_stats,
            authenticator,
            obfuscator,
        })
    }

    pub async fn serve(&self) -> anyhow::Result<()> {
        info!("Hysteria server starting on {}", self.config.listen);

        // Load TLS configuration
        let tls_config = self.load_tls_config()?;
        
        // Create QUIC server config
        let mut server_config = QuinnServerConfig::with_crypto(Arc::new(tls_config));
        let transport_config = hysteria_core::transport::QuicConfig::default().to_quinn_config();
        server_config.transport_config(Arc::new(transport_config));

        // Create QUIC endpoint
        let endpoint = Endpoint::server(server_config, self.config.listen)?;
        info!("Server is ready to accept connections");

        // Accept connections
        loop {
            tokio::select! {
                Some(conn) = endpoint.accept() => {
                    let server = self.clone_arc();
                    tokio::spawn(async move {
                        if let Err(e) = server.handle_connection(conn).await {
                            error!("Connection error: {}", e);
                        }
                    });
                }
                _ = tokio::signal::ctrl_c() => {
                    info!("Shutting down server");
                    break;
                }
            }
        }

        Ok(())
    }

    fn clone_arc(&self) -> Arc<Self> {
        Arc::new(Self {
            config: Arc::clone(&self.config),
            traffic_stats: Arc::clone(&self.traffic_stats),
            authenticator: Arc::clone(&self.authenticator),
            obfuscator: self.obfuscator.as_ref().map(Arc::clone),
        })
    }

    async fn handle_connection(&self, connecting: quinn::Connecting) -> anyhow::Result<()> {
        let remote_addr = connecting.remote_address();
        debug!("New connection from {}", remote_addr);

        let connection = connecting.await?;
        info!("Connection established from {}", remote_addr);

        // Handle HTTP/3 authentication
        let (mut send, mut recv) = connection.accept_bi().await?;
        
        // Read authentication request
        let mut buf = vec![0u8; 4096];
        let n = recv.read(&mut buf).await?.unwrap_or(0);
        
        if n == 0 {
            warn!("Empty authentication request from {}", remote_addr);
            return Ok(());
        }

        // Parse auth request (simplified)
        let auth_data = String::from_utf8_lossy(&buf[..n]);
        let auth_req = AuthRequest::new(auth_data.to_string(), 0);

        // Authenticate
        let (ok, user_id) = self.handle_auth_request(&auth_req, remote_addr).await?;
        
        if !ok {
            warn!("Authentication failed for {}", remote_addr);
            send.finish().await?;
            return Ok(());
        }

        info!("User {} authenticated from {}", user_id, remote_addr);
        
        // Send auth response
        let auth_resp = self.create_auth_response(!self.config.disable_udp);
        let resp_data = serde_json::to_vec(&auth_resp)?;
        send.write_all(&resp_data).await?;
        send.finish().await?;

        // Update traffic stats
        self.traffic_stats.set_online(&user_id, true).await;

        // Handle proxy requests
        while let Ok((mut send_stream, mut recv_stream)) = connection.accept_bi().await {
            let server = Arc::new(self.clone_arc());
            let user_id_clone = user_id.clone();
            
            tokio::spawn(async move {
                if let Err(e) = server.handle_stream(send_stream, recv_stream, &user_id_clone).await {
                    error!("Stream error: {}", e);
                }
            });
        }

        self.traffic_stats.set_online(&user_id, false).await;
        info!("Connection closed for user {}", user_id);

        Ok(())
    }

    async fn handle_stream(
        &self,
        mut send: quinn::SendStream,
        mut recv: quinn::RecvStream,
        user_id: &str,
    ) -> anyhow::Result<()> {
        // Read TCP request
        let mut buf = vec![0u8; 4096];
        let n = recv.read(&mut buf).await?.unwrap_or(0);
        
        if n == 0 {
            return Ok(());
        }

        let tcp_req = TcpRequest::decode(&buf[..n])?;
        debug!("TCP request to {} from user {}", tcp_req.address, user_id);

        // Connect to target
        match TcpStream::connect(&tcp_req.address).await {
            Ok(mut target) => {
                // Send success response
                let resp = TcpResponse::ok();
                send.write_all(&resp.encode()).await?;

                // Relay data
                let (mut target_read, mut target_write) = target.split();
                
                let upload = async {
                    let mut buf = vec![0u8; 8192];
                    let mut total = 0u64;
                    loop {
                        match recv.read(&mut buf).await {
                            Ok(Some(n)) => {
                                target_write.write_all(&buf[..n]).await?;
                                total += n as u64;
                            }
                            Ok(None) => break,
                            Err(e) => return Err(e.into()),
                        }
                    }
                    Ok::<u64, anyhow::Error>(total)
                };

                let download = async {
                    let mut buf = vec![0u8; 8192];
                    let mut total = 0u64;
                    loop {
                        match target_read.read(&mut buf).await {
                            Ok(0) => break,
                            Ok(n) => {
                                send.write_all(&buf[..n]).await?;
                                total += n as u64;
                            }
                            Err(e) => return Err(e.into()),
                        }
                    }
                    Ok::<u64, anyhow::Error>(total)
                };

                let (tx_result, rx_result) = tokio::join!(upload, download);
                let tx = tx_result.unwrap_or(0);
                let rx = rx_result.unwrap_or(0);

                // Log traffic
                self.traffic_stats.log_traffic(user_id, tx, rx).await;
                debug!("Stream finished: tx={}, rx={}", tx, rx);
            }
            Err(e) => {
                warn!("Failed to connect to {}: {}", tcp_req.address, e);
                let resp = TcpResponse::error(format!("Connection failed: {}", e));
                send.write_all(&resp.encode()).await?;
            }
        }

        send.finish().await?;
        Ok(())
    }

    fn load_tls_config(&self) -> anyhow::Result<RustlsServerConfig> {
        let tls_config = self.config.tls.as_ref()
            .ok_or_else(|| anyhow::anyhow!("TLS configuration required"))?;

        // Load certificates
        let cert_file = std::fs::File::open(&tls_config.cert)?;
        let mut cert_reader = std::io::BufReader::new(cert_file);
        let certs = rustls_pemfile::certs(&mut cert_reader)
            .collect::<Result<Vec<_>, _>>()?;

        // Load private key
        let key_file = std::fs::File::open(&tls_config.key)?;
        let mut key_reader = std::io::BufReader::new(key_file);
        let key = rustls_pemfile::private_key(&mut key_reader)?
            .ok_or_else(|| anyhow::anyhow!("No private key found"))?;

        let mut config = RustlsServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(certs, key)?;

        config.alpn_protocols = vec![b"h3".to_vec()];

        Ok(config)
    }

    pub fn get_traffic_stats(&self) -> Arc<TrafficStats> {
        Arc::clone(&self.traffic_stats)
    }

    async fn handle_auth_request(
        &self,
        auth_req: &AuthRequest,
        remote_addr: std::net::SocketAddr,
    ) -> anyhow::Result<(bool, String)> {
        self.authenticator
            .authenticate(remote_addr, &auth_req.auth, auth_req.rx)
            .await
    }

    fn create_auth_response(&self, udp_enabled: bool) -> AuthResponse {
        let (up, down) = self.config.parse_bandwidth();
        
        AuthResponse {
            udp_enabled,
            rx: down,
            rx_auto: self.config.ignore_client_bandwidth,
            padding: String::new(),
        }
    }
}
